{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Ch 5 - Actor-Critic Models\n",
    "### Deep Reinforcement Learning in Action"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Listing 5.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import multiprocessing as mp\n",
    "import numpy as np\n",
    "def square(x): #A\n",
    "    return np.square(x)\n",
    "x = np.arange(64) #B\n",
    "print(x)\n",
    "print(mp.cpu_count())\n",
    "pool = mp.Pool(8) #C\n",
    "squared = pool.map(square, [x[8*i:8*i+8] for i in range(8)])\n",
    "print(squared)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Listing 5.2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def square(i, x, queue):\n",
    "    print(\"In process {}\".format(i,))\n",
    "    queue.put(np.square(x))\n",
    "processes = [] #A\n",
    "queue = mp.Queue() #B\n",
    "x = np.arange(64) #C\n",
    "for i in range(8): #D\n",
    "    start_index = 8*i\n",
    "    proc = mp.Process(target=square,args=(i,x[start_index:start_index+8], queue)) \n",
    "    proc.start()\n",
    "    processes.append(proc)\n",
    "    \n",
    "for proc in processes: #E\n",
    "    proc.join()\n",
    "    \n",
    "for proc in processes: #F\n",
    "    proc.terminate()\n",
    "results = []\n",
    "while not queue.empty(): #G\n",
    "    results.append(queue.get())\n",
    "\n",
    "print(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Listing 5.4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch import optim\n",
    "import numpy as np\n",
    "from torch.nn import functional as F\n",
    "import gym\n",
    "import torch.multiprocessing as mp #A\n",
    "\n",
    "class ActorCritic(nn.Module): #B\n",
    "    def __init__(self):\n",
    "        super(ActorCritic, self).__init__()\n",
    "        self.l1 = nn.Linear(4,25)\n",
    "        self.l2 = nn.Linear(25,50)\n",
    "        self.actor_lin1 = nn.Linear(50,2)\n",
    "        self.l3 = nn.Linear(50,25)\n",
    "        self.critic_lin1 = nn.Linear(25,1)\n",
    "    def forward(self,x):\n",
    "        x = F.normalize(x,dim=0)\n",
    "        y = F.relu(self.l1(x))\n",
    "        y = F.relu(self.l2(y))\n",
    "        actor = F.log_softmax(self.actor_lin1(y),dim=0) #C\n",
    "        c = F.relu(self.l3(y.detach()))\n",
    "        critic = torch.tanh(self.critic_lin1(c)) #D\n",
    "        return actor, critic #E"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Listing 5.6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def worker(t, worker_model, counter, params):\n",
    "    worker_env = gym.make(\"CartPole-v1\")\n",
    "    worker_env.reset()\n",
    "    worker_opt = optim.Adam(lr=1e-4,params=worker_model.parameters()) #A\n",
    "    worker_opt.zero_grad()\n",
    "    for i in range(params['epochs']):\n",
    "        worker_opt.zero_grad()\n",
    "        values, logprobs, rewards = run_episode(worker_env,worker_model) #B \n",
    "        actor_loss,critic_loss,eplen = update_params(worker_opt,values,logprobs,rewards) #C\n",
    "        counter.value = counter.value + 1 #D"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Listing 5.7"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_episode(worker_env, worker_model):\n",
    "    state = torch.from_numpy(worker_env.env.state).float() #A\n",
    "    values, logprobs, rewards = [],[],[] #B\n",
    "    done = False\n",
    "    j=0\n",
    "    while (done == False): #C\n",
    "        j+=1\n",
    "        policy, value = worker_model(state) #D\n",
    "        values.append(value)\n",
    "        logits = policy.view(-1)\n",
    "        action_dist = torch.distributions.Categorical(logits=logits)\n",
    "        action = action_dist.sample() #E\n",
    "        logprob_ = policy.view(-1)[action]\n",
    "        logprobs.append(logprob_)\n",
    "        state_, _, done, info = worker_env.step(action.detach().numpy())\n",
    "        state = torch.from_numpy(state_).float()\n",
    "        if done: #F\n",
    "            reward = -10\n",
    "            worker_env.reset()\n",
    "        else:\n",
    "            reward = 1.0\n",
    "        rewards.append(reward)\n",
    "    return values, logprobs, rewards"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Listing 5.8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def update_params(worker_opt,values,logprobs,rewards,clc=0.1,gamma=0.95):\n",
    "        rewards = torch.Tensor(rewards).flip(dims=(0,)).view(-1) #A\n",
    "        logprobs = torch.stack(logprobs).flip(dims=(0,)).view(-1)\n",
    "        values = torch.stack(values).flip(dims=(0,)).view(-1)\n",
    "        Returns = []\n",
    "        ret_ = torch.Tensor([0])\n",
    "        for r in range(rewards.shape[0]): #B\n",
    "            ret_ = rewards[r] + gamma * ret_\n",
    "            Returns.append(ret_)\n",
    "        Returns = torch.stack(Returns).view(-1)\n",
    "        Returns = F.normalize(Returns,dim=0)\n",
    "        actor_loss = -1*logprobs * (Returns - values.detach()) #C\n",
    "        critic_loss = torch.pow(values - Returns,2) #D\n",
    "        loss = actor_loss.sum() + clc*critic_loss.sum() #E\n",
    "        loss.backward()\n",
    "        worker_opt.step()\n",
    "        return actor_loss, critic_loss, len(rewards)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Listing 5.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MasterNode = ActorCritic() #A\n",
    "MasterNode.share_memory() #B\n",
    "processes = [] #C\n",
    "params = {\n",
    "    'epochs':1000,\n",
    "    'n_workers':7,\n",
    "}\n",
    "counter = mp.Value('i',0) #D\n",
    "for i in range(params['n_workers']):\n",
    "    p = mp.Process(target=worker, args=(i,MasterNode,counter,params)) #E\n",
    "    p.start() \n",
    "    processes.append(p)\n",
    "for p in processes: #F\n",
    "    p.join()\n",
    "for p in processes: #G\n",
    "    p.terminate()\n",
    "    \n",
    "print(counter.value,processes[1].exitcode) #H"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Test the trained agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.make(\"CartPole-v1\")\n",
    "env.reset()\n",
    "\n",
    "for i in range(100):\n",
    "    state_ = np.array(env.env.state)\n",
    "    state = torch.from_numpy(state_).float()\n",
    "    logits,value = MasterNode(state)\n",
    "    action_dist = torch.distributions.Categorical(logits=logits)\n",
    "    action = action_dist.sample()\n",
    "    state2, reward, done, info = env.step(action.detach().numpy())\n",
    "    if done:\n",
    "        print(\"Lost\")\n",
    "        env.reset()\n",
    "    state_ = np.array(env.env.state)\n",
    "    state = torch.from_numpy(state_).float()\n",
    "    env.render()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Listing 5.9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_episode(worker_env, worker_model, N_steps=10):\n",
    "    raw_state = np.array(worker_env.env.state)\n",
    "    state = torch.from_numpy(raw_state).float()\n",
    "    values, logprobs, rewards = [],[],[]\n",
    "    done = False\n",
    "    j=0\n",
    "    G=torch.Tensor([0]) #A\n",
    "    while (j < N_steps and done == False): #B\n",
    "        j+=1\n",
    "        policy, value = worker_model(state)\n",
    "        values.append(value)\n",
    "        logits = policy.view(-1)\n",
    "        action_dist = torch.distributions.Categorical(logits=logits)\n",
    "        action = action_dist.sample()\n",
    "        logprob_ = policy.view(-1)[action]\n",
    "        logprobs.append(logprob_)\n",
    "        state_, _, done, info = worker_env.step(action.detach().numpy())\n",
    "        state = torch.from_numpy(state_).float()\n",
    "        if done:\n",
    "            reward = -10\n",
    "            worker_env.reset()\n",
    "        else: #C\n",
    "            reward = 1.0\n",
    "            G = value.detach()\n",
    "        rewards.append(reward)\n",
    "    return values, logprobs, rewards, G"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Listing 5.10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Simulated rewards for 3 steps\n",
    "r1 = [1,1,-1]\n",
    "r2 = [1,1,1]\n",
    "R1,R2 = 0.0,0.0\n",
    "#No bootstrapping\n",
    "for i in range(len(r1)-1,0,-1): \n",
    "    R1 = r1[i] + 0.99*R1\n",
    "for i in range(len(r2)-1,0,-1):\n",
    "    R2 = r2[i] + 0.99*R2\n",
    "print(\"No bootstrapping\")\n",
    "print(R1,R2)\n",
    "#With bootstrapping\n",
    "R1,R2 = 1.0,1.0\n",
    "for i in range(len(r1)-1,0,-1):\n",
    "    R1 = r1[i] + 0.99*R1\n",
    "for i in range(len(r2)-1,0,-1):\n",
    "    R2 = r2[i] + 0.99*R2\n",
    "print(\"With bootstrapping\")\n",
    "print(R1,R2)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:deeprl]",
   "language": "python",
   "name": "conda-env-deeprl-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
